package com.ai.common.utils;

import java.util.HashMap;
import java.util.UUID;

public class TraceContext {

    private UUID uuid = java.util.UUID.randomUUID();

    private static final ThreadLocal<TraceContext> threadLocalContext = new ThreadLocal<TraceContext>() {
        /**
         * ThreadLocal没有被当前线程赋值时或当前线程刚调用remove方法后调用get方法，返回此方法值
         */
        @Override
        protected TraceContext initialValue() {
            return new TraceContext();
        }
    };

    public static TraceContext getCurrent() {
        if (threadLocalContext.get() == null)
            threadLocalContext.set(new TraceContext());
        return threadLocalContext.get();
    }

    public static void setCurrent(TraceContext context) {
        threadLocalContext.set(context);
    }

    public static void clearCurrent() {
        threadLocalContext.remove();
    }



    public UUID getUuid() {
        return this.uuid;
    }

    public void addTags(String key, String value) {
        if (tags == null)
            tags = new HashMap<String, String>();
        if (tags.size() > 50)
            tags = new HashMap<String, String>();
        tags.put(key, value);
    }

    public void clearTags() {
        if (tags == null)
            return;
        tags.clear();
    }

    public HashMap<String, String> getTags() {
        if (tags == null)
            tags = new HashMap<String, String>();
        return tags;
    }

    public <T> T getTag(String key, Class<T> cls) {
        return getTag(key, cls, null);
    }

    public <T> T getTag(String key, Class<T> cls, T defaultValue) {
        if (key == null) return defaultValue;
        Object item = tags.get(key);
        if (item != null && item.getClass().equals(cls))
            return (T) item;
        else
            return defaultValue;
    }

    private HashMap<String, String> tags = null;

    private HashMap<String, Object> objTags = null;

    public void addObjTags(String key, Object value) {
        if (objTags == null)
            objTags = new HashMap<String, Object>();
        if (objTags.size() > 50)
            objTags = new HashMap<String, Object>();
        objTags.put(key, value);
    }

    public void clearObjTags() {
        if (objTags == null)
            return;
        objTags.clear();
    }

    public HashMap<String, Object> getObjTags() {
        if (objTags == null)
            objTags = new HashMap<String, Object>();
        return objTags;
    }

    public <T> T getObjTag(String key, Class<T> cls) {
        return getObjTag(key, cls, null);
    }

    public <T> T getObjTag(String key, Class<T> cls, T defaultValue) {
        if (key == null) return defaultValue;
        Object item = objTags.get(key);
        if (item != null && item.getClass().equals(cls))
            return (T) item;
        else
            return defaultValue;
    }
}
